"""
优化器钩子模块，用于监控优化器的状态。
"""
from typing import Callable, Optional, Any, Dict
import torch

class OptimizerHook:
    """优化器钩子"""
    
    def __init__(self,
                 name: str,
                 optimizer: Any,
                 callback: Callable[[str, Dict[str, Any], float], None]):
        """
        初始化优化器钩子
        
        Args:
            name: 优化器名称
            optimizer: 优化器实例
            callback: 回调函数
        """
        self.name = name
        self.optimizer = optimizer
        self.callback = callback
        self.hook = None
        self.register()
    
    def register(self) -> None:
        """注册钩子"""
        if self.hook is not None:
            return
        
        def hook(optimizer: Any, state: Dict[str, Any], group: Dict[str, Any]) -> None:
            """优化器钩子函数"""
            # 获取学习率
            lr = group.get('lr', 0.0)
            
            # 调用回调函数
            self.callback(self.name, state, lr)
        
        # 注册钩子
        if hasattr(self.optimizer, 'register_step_post_hook'):
            self.hook = self.optimizer.register_step_post_hook(hook)
        else:
            # 对于不支持钩子的优化器，使用包装器
            original_step = self.optimizer.step
            def wrapped_step(*args, **kwargs):
                result = original_step(*args, **kwargs)
                # 获取第一个参数组的状态
                if self.optimizer.param_groups:
                    state = self.optimizer.state[self.optimizer.param_groups[0]['params'][0]]
                    group = self.optimizer.param_groups[0]
                    hook(self.optimizer, state, group)
                return result
            self.optimizer.step = wrapped_step
    
    def remove(self) -> None:
        """移除钩子"""
        if self.hook is not None:
            self.hook.remove()
            self.hook = None
        else:
            # 恢复原始step方法
            if hasattr(self.optimizer, '_original_step'):
                self.optimizer.step = self.optimizer._original_step
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.remove() 